import os
from langchain_core.pydantic_v1 import BaseModel, Field, validator


class PythonScriptFileModel(BaseModel):
    python_script: str = Field(description="A Python script to generate fake data, and the method signature for generating the fake data file is fixed as: generate_fake_data_file(data_num = 1)")
    script_filename: str = Field(description="The filename for the Python script to generate fake data should be adjusted based on the content of [Example Data]")
    data_type: str = Field(description="The data type of the [Example Data], which can be 'json' or 'sql'")

    @validator("python_script")
    def python_script_in_code_block(cls, field):
        if "from seed_data_seeker" not in field and "import seed_data_seeker" not in field:
            raise ValueError("Badly formed python script!")
        if "from seed_data_seeker" in field:
            new_field = field.replace('from seed_data_seeker', 'from fake_data_script_executer')
        elif "import seed_data_seeker" in field:
            new_field = field.replace('seed_data_seeker', 'fake_data_script_executer')
        return new_field
    
    @validator("script_filename")
    def python_filename_ends_with_py(cls, field):
        _, extname = os.path.splitext(field)
        if extname != ".py":
            raise ValueError("Badly script filename!")
        return field
    
    @validator("data_type")
    def data_type_in_enum(cls, field):
        if field not in ["json", "sql"]:
            raise ValueError("Invalid data type!")
        return field
